package agent

import (
	"context"
	"fmt"
	"log"
	"time"

	"demo2/server/storage"
)

// Agent 智能代理核心
type Agent struct {
	memory     *Memory
	planner    *Planner
	executor   *Executor
	tools      *ToolManager
	llm        LLMInterface
	store      *storage.RedisStore
}

// AgentConfig Agent配置
type AgentConfig struct {
	MaxMemorySize    int
	PlanningTimeout  time.Duration
	ExecutionTimeout time.Duration
}

// NewAgent 创建新的Agent实例
func NewAgent(llm LLMInterface, store *storage.RedisStore, config AgentConfig) *Agent {
	memory := NewMemory(config.MaxMemorySize)
	tools := NewToolManager()
	planner := NewPlanner(llm, tools, config.PlanningTimeout)
	executor := NewExecutor(tools, config.ExecutionTimeout)

	agent := &Agent{
		memory:   memory,
		planner:  planner,
		executor: executor,
		tools:    tools,
		llm:      llm,
		store:    store,
	}

	// 注册默认工具
	agent.registerDefaultTools()

	return agent
}

// registerDefaultTools 注册默认工具
func (a *Agent) registerDefaultTools() {
	// 时间查询工具
	a.tools.RegisterTool(&Tool{
		Name:        "get_time",
		Description: "查询指定地点的当前时间",
		Parameters: map[string]interface{}{
			"location": map[string]interface{}{
				"type":        "string",
				"description": "地点名称，如：北京、纽约、东京等",
				"required":    true,
			},
		},
		Handler: a.handleTimeQuery,
	})

	// 计算工具
	a.tools.RegisterTool(&Tool{
		Name:        "calculate",
		Description: "执行数学计算",
		Parameters: map[string]interface{}{
			"expression": map[string]interface{}{
				"type":        "string",
				"description": "数学表达式，支持+、-、*、/、sqrt()、pow()等",
				"required":    true,
			},
		},
		Handler: a.handleCalculation,
	})

	// 记忆查询工具
	a.tools.RegisterTool(&Tool{
		Name:        "search_memory",
		Description: "搜索历史对话记忆",
		Parameters: map[string]interface{}{
			"query": map[string]interface{}{
				"type":        "string",
				"description": "搜索关键词",
				"required":    true,
			},
		},
		Handler: a.handleMemorySearch,
	})
}

// ProcessMessage 处理用户消息
func (a *Agent) ProcessMessage(ctx context.Context, sessionID, userMessage string) (*AgentResponse, error) {
	log.Printf("[AGENT] 🤖 开始处理消息，会话ID: %s", sessionID)

	// 1. 加载会话记忆
	if err := a.loadSessionMemory(ctx, sessionID); err != nil {
		log.Printf("[AGENT] ❌ 加载会话记忆失败: %v", err)
		return nil, err
	}

	// 2. 添加用户消息到记忆
	a.memory.AddMessage(Message{
		Role:      "user",
		Content:   userMessage,
		Timestamp: time.Now(),
		SessionID: sessionID,
	})

	// 3. 分析用户意图并制定计划
	plan, err := a.planner.CreatePlan(ctx, userMessage, a.memory.GetRecentMessages(10))
	if err != nil {
		log.Printf("[AGENT] ❌ 制定计划失败: %v", err)
		return nil, err
	}

	log.Printf("[AGENT] 📋 制定计划完成，步骤数: %d", len(plan.Steps))

	// 4. 执行计划
	result, err := a.executor.ExecutePlan(ctx, plan)
	if err != nil {
		log.Printf("[AGENT] ❌ 执行计划失败: %v", err)
		return nil, err
	}

	// 5. 生成最终回复
	response, err := a.generateResponse(ctx, userMessage, result)
	if err != nil {
		log.Printf("[AGENT] ❌ 生成回复失败: %v", err)
		return nil, err
	}

	// 6. 添加回复到记忆
	a.memory.AddMessage(Message{
		Role:      "assistant",
		Content:   response.Content,
		Timestamp: time.Now(),
		SessionID: sessionID,
	})

	// 7. 保存记忆到存储
	if err := a.saveSessionMemory(ctx, sessionID); err != nil {
		log.Printf("[AGENT] ⚠️ 保存会话记忆失败: %v", err)
	}

	log.Printf("[AGENT] ✅ 消息处理完成")
	return response, nil
}

// ProcessStreamMessage 处理流式消息
func (a *Agent) ProcessStreamMessage(ctx context.Context, sessionID, userMessage string, responseChan chan<- string) error {
	defer close(responseChan)

	// 加载会话记忆
	if err := a.loadSessionMemory(ctx, sessionID); err != nil {
		responseChan <- fmt.Sprintf("加载会话记忆失败: %v", err)
		return err
	}

	// 添加用户消息到记忆
	a.memory.AddMessage(Message{
		Role:      "user",
		Content:   userMessage,
		Timestamp: time.Now(),
		SessionID: sessionID,
	})

	// 制定计划
	plan, err := a.planner.CreatePlan(ctx, userMessage, a.memory.GetRecentMessages(10))
	if err != nil {
		responseChan <- fmt.Sprintf("制定计划失败: %v", err)
		return err
	}

	// 执行计划
	result, err := a.executor.ExecutePlan(ctx, plan)
	if err != nil {
		responseChan <- fmt.Sprintf("执行计划失败: %v", err)
		return err
	}

	// 生成流式回复
	fullResponse := ""
	streamChan := make(chan string, 100)
	
	go func() {
		defer close(streamChan)
		if err := a.generateStreamResponse(ctx, userMessage, result, streamChan); err != nil {
			streamChan <- fmt.Sprintf("生成回复失败: %v", err)
		}
	}()

	// 转发流式响应
	for chunk := range streamChan {
		fullResponse += chunk
		select {
		case <-ctx.Done():
			return ctx.Err()
		case responseChan <- chunk:
		}
	}

	// 保存完整回复到记忆
	if fullResponse != "" {
		a.memory.AddMessage(Message{
			Role:      "assistant",
			Content:   fullResponse,
			Timestamp: time.Now(),
			SessionID: sessionID,
		})

		// 异步保存记忆
		go func() {
			a.saveSessionMemory(context.Background(), sessionID)
		}()
	}

	return nil
}

// loadSessionMemory 加载会话记忆
func (a *Agent) loadSessionMemory(ctx context.Context, sessionID string) error {
	messages, err := a.store.GetMessages(ctx, sessionID)
	if err != nil {
		return err
	}

	for _, msg := range messages {
		a.memory.AddMessage(Message{
			Role:      msg.Role,
			Content:   msg.Content,
			Timestamp: msg.Timestamp,
			SessionID: sessionID,
		})
	}

	return nil
}

// saveSessionMemory 保存会话记忆
func (a *Agent) saveSessionMemory(ctx context.Context, sessionID string) error {
	messages := a.memory.GetSessionMessages(sessionID)
	
	for _, msg := range messages {
		storageMsg := &storage.Message{
			ID:        fmt.Sprintf("%s_%d", sessionID, msg.Timestamp.UnixNano()),
			SessionID: sessionID,
			Role:      msg.Role,
			Content:   msg.Content,
			Timestamp: msg.Timestamp,
		}
		
		if err := a.store.SaveMessage(ctx, storageMsg); err != nil {
			return err
		}
	}

	return nil
}

// generateResponse 生成最终回复
func (a *Agent) generateResponse(ctx context.Context, userMessage string, executionResult *ExecutionResult) (*AgentResponse, error) {
	// 构建系统提示
	systemPrompt := a.buildSystemPrompt(executionResult)
	
	// 构建消息历史
	messages := []LLMMessage{
		{Role: "system", Content: systemPrompt},
	}
	
	// 添加最近的对话历史
	recentMessages := a.memory.GetRecentMessages(5)
	for _, msg := range recentMessages {
		messages = append(messages, LLMMessage{
			Role:    msg.Role,
			Content: msg.Content,
		})
	}

	// 调用LLM生成回复
	response, err := a.llm.GenerateResponse(ctx, messages)
	if err != nil {
		return nil, err
	}

	return &AgentResponse{
		Content:         response,
		ExecutionResult: executionResult,
		Timestamp:       time.Now(),
	}, nil
}

// generateStreamResponse 生成流式回复
func (a *Agent) generateStreamResponse(ctx context.Context, userMessage string, executionResult *ExecutionResult, responseChan chan<- string) error {
	// 如果有工具执行结果，直接返回
	if executionResult != nil && len(executionResult.ToolResults) > 0 {
		for _, result := range executionResult.ToolResults {
			if result.Success {
				responseChan <- result.Output
				return nil
			}
		}
	}

	// 构建系统提示
	systemPrompt := a.buildSystemPrompt(executionResult)
	
	// 构建消息历史
	messages := []LLMMessage{
		{Role: "system", Content: systemPrompt},
	}
	
	// 添加最近的对话历史
	recentMessages := a.memory.GetRecentMessages(5)
	for _, msg := range recentMessages {
		messages = append(messages, LLMMessage{
			Role:    msg.Role,
			Content: msg.Content,
		})
	}

	// 调用LLM生成流式回复
	return a.llm.GenerateStreamResponse(ctx, messages, responseChan)
}

// buildSystemPrompt 构建系统提示
func (a *Agent) buildSystemPrompt(executionResult *ExecutionResult) string {
	prompt := `你是一个智能助手，具有以下能力：

1. 时间查询 - 可以查询世界各地的当前时间
2. 数学计算 - 可以进行各种数学运算
3. 记忆搜索 - 可以搜索历史对话内容

你的回答应该：
- 准确、有用、友好
- 基于工具执行结果提供信息
- 保持对话的连贯性
- 适当使用上下文信息`

	if executionResult != nil && len(executionResult.ToolResults) > 0 {
		prompt += "\n\n工具执行结果：\n"
		for _, result := range executionResult.ToolResults {
			if result.Success {
				prompt += fmt.Sprintf("- %s: %s\n", result.ToolName, result.Output)
			} else {
				prompt += fmt.Sprintf("- %s: 执行失败 - %s\n", result.ToolName, result.Error)
			}
		}
	}

	return prompt
}

// GetMemory 获取记忆管理器
func (a *Agent) GetMemory() *Memory {
	return a.memory
}

// GetTools 获取工具管理器
func (a *Agent) GetTools() *ToolManager {
	return a.tools
}

// Close 关闭Agent
func (a *Agent) Close() error {
	return a.executor.Close()
}
